"""
Used to extract the R,L,M from two port measurement results.

+V1-I1->|           |-<--I2---V2+
      |--|         |--|
      |X1|   <-M-> |X2|
      |--|         |--|
        |           | 
---------           ------------
Port1(Z0)                 Port2(Z0)
===========================================
1. terminate p2 with z0
let A = V1/I1 = x1 + w^2*m^2/(x2+z0)
and B = V1/I2 = j*(z0+x2+w^2*m^2/x1)/(w*m/x1)
2. terminate p1 with z0
let C = V2/I2 = x2 + w^2*m^2/(x1+z0)
and D = V2/I1 = j*(z0+x1+w^2*m^2/x2)/(w*m/x2)
A --dual-- C    B --dual-- D  1 --dual-- 2
3. From S-parameters,
A = z0*(1+s11)/(1-s11) ## v1 = sqrt(z0)*(a1+b1) and b1 = s11*a1 since a2=0 when port 2 terminated with z0
                       ## i1 = (a1-b1)/sqrt(z0)
B = z0*(1+s11)/(-s21)  ## i2 = (0-b2)/sqrt(z0)  and b2 = s21*a1
                       ## v2 = sqrt(z0)*(0+b2)
C = z0*(1+s22)/(1-s22)
D = z0*(1+s22)/(-s12)
4. Calculate ABCD from S-parameters first
5.  From ABCD
x1 = A*(BD+C*z0)/(BD-AC) (x1 could be inductor, so x1=R1+jwL1)
x2 = C*(BD+A*z0)/(BD-AC) (or anything with capacitor or resistor)
jwm = BC*(A+z0)/(AC-BD)
"""


from cmath import *
from numpy import array, linspace, logspace, pi, log, sin, log10, sqrt, vstack, hstack, transpose, abs
from scipy.interpolate import splprep, splrep, splev
from pychartdir import *
import os,re

j= complex(0,1)
mu = 4.0e-7*pi
ep = 8.854e-12

from numpy import exp

def pol(a,b=None):
    "generate complex number from (mag,deg) form"
    if type(a)==type(list()):
        return a[0]*exp(complex(0,1)*a[1]*pi/180)
    return a*exp(complex(0,1)*b*pi/180)

def arg(a):
    "return the argument(in degree) of complex number"
    return ((log(a/abs(a))/complex(0,1)).real)*180/pi

def aa(a):
    "return both the amplitude and the argument"
    return [abs(a), arg(a)]

def abcd(s11,s12,s21,s22,z0=50):
    "return ABCD from S-parameters"
    try:
        return (1+s11)*z0/(1-s11),(1+s11)*z0/(-s21),(1+s22)*z0/(1-s22),(1+s22)*z0/(-s12)
    except ZeroDivisionError:
        return (1+s11)*z0/(1-s11+1e-21),(1+s11)*z0/(-s21+1e-21),(1+s22)*z0/(1-s22+1e-21),(1+s22)*z0/(-s12+1e-21)

def rlm(a,b,c,d,z0=50):
    "return x1,x2,jwm from ABCD"
    return a*(b*d+c*z0)/(b*d-a*c),c*(b*d+a*z0)/(b*d-a*c),b*c*(a+z0)/(a*c-b*d)

freqpt = re.compile(r"^# (.*?) .* ([.0-9]+)")
class Cord:
    class Axis:
        def __init__(self):
            self.mode=False
            self.min=None
            self.max=None
    def __init__(self):
        self.xaxis = self.Axis()
        self.yaxis = self.Axis()
    def setXmin(self,x):
        self.xaxis.min=x
        return self
    def setXmax(self,x):
        self.xaxis.max=x
        return self
    def setYmin(self,y):
        self.yaxis.min=y
        return self
    def setYmax(self,y):
        self.yaxis.max=y
        return self
    def setXMode(self, mode=False):
        if mode: self.xaxis.mode=True
        else: self.xaxis.mode=False
        return self
    def setYMode(self, mode=False):
        if mode: self.yaxis.mode=True
        else: self.yaxis.mode=False
        return self
    def __str__(self):
        return "x (%s:%s) log_OnOff:%s\ty (%s:%s) log_OnOff:%s"%(self.xaxis.min,
                                                                 self.xaxis.max,
                                                                 self.xaxis.mode,
                                                                 self.yaxis.min,
                                                                 self.yaxis.max,
                                                                 self.yaxis.mode)

def setAxis(c,cord):
    if None in [cord.xaxis.min,cord.xaxis.max]:
        if cord.xaxis.mode:
            c.xAxis().setLogScale3()
        else:
            c.xAxis().setLinearScale3()
    else:
        if cord.xaxis.mode:
            c.xAxis().setLogScale(cord.xaxis.min,cord.xaxis.max)
        else:
            c.xAxis().setLinearScale(cord.xaxis.min,cord.xaxis.max)    
    if None in [cord.yaxis.min,cord.yaxis.max]:
        if cord.yaxis.mode:
            c.yAxis().setLogScale3()
        else:
            c.yAxis().setLinearScale3()
    else:
        if cord.yaxis.mode:
            c.yAxis().setLogScale(cord.yaxis.min,cord.yaxis.max)
        else:
            c.yAxis().setLinearScale(cord.yaxis.min,cord.yaxis.max) 
AutoCord = Cord()
LogXCord = Cord()
LogXCord.setXMode(True)
LogYCord = Cord()
LogYCord.setYMode(True)
LogLogCord = Cord()
LogLogCord.setXMode(True).setYMode(True)
CurCord = Cord()


def sp(name,AmpArgMode=False):
    """read S parameters from s2p file"""
    try:
        fin = open(name, 'rb')
        dt = fin.read()
        fin.close()
    except:
        print "file '%s' not found in path %s\nUse direct input...\n"%(name, os.getcwd())
        dt = raw_input()
        name = "rawinput_"
    lines = dt.splitlines()
    scale,z0 = 1.0,50.0
    mat = []
    for i in range(len(lines)):
        line=lines[i].strip()
        if line.startswith("!"):
            continue
        if line.startswith("#"):
            freqmt = freqpt.findall(line)
            if freqmt:
                freqtext,z0=freqmt[0]
                z0 = float(z0)
                if freqtext.upper().startswith("K"):
                    scale=1e3
                elif freqtext.upper().startswith("M"):
                    scale=1e6
                elif freqtext.upper().startswith("G"):                    
                    scale=1e9
                elif freqtext.upper().startswith("T"):
                    scale=1e12
                print "frequency in (%E)Hz\n"%scale
            continue
        data=[float(x) for x in line.split()]
        freq, data = data[0]*scale,data[1:]
        if len(data)!=8:
            print 'wrong line: %s'%line
            continue
        else:
            if AmpArgMode:
                s11,s12,s21,s22 = pol(data[0],data[1]),pol(data[2],data[3]),pol(data[4],data[5]),pol(data[6],data[7])
            else:
                s11,s12,s21,s22 = complex(data[0],data[1]),complex(data[2],data[3]),complex(data[4],data[5]),complex(data[6],data[7])
        mat.append((freq,s11,s12,s21,s22))
    return array(mat),name    

def read_s2p(name, bTwoPorts =True,AmpArgMode=False):
    """if AmpArgMode is True, then use pol(amp,arg) to generate complex number
    otherwise, use complex(real,imag)
    """
    try:
        fin = open(name, 'rb')
        dt = fin.read()
        fin.close()
    except:
        print "file '%s' not found in path %s\nUse direct input...\n"%(name, os.getcwd())
        dt = raw_input()
        name = "rawinput_"
    lines = dt.splitlines()
    scale,z0 = 1.0,50.0
    mat = []
    for i in range(len(lines)):
        line=lines[i].strip()
        if line.startswith("!"):
            continue
        if line.startswith("#"):
            freqmt = freqpt.findall(line)
            if freqmt:
                freqtext,z0=freqmt[0]
                z0 = float(z0)
                if freqtext.upper().startswith("K"):
                    scale=1e3
                elif freqtext.upper().startswith("M"):
                    scale=1e6
                elif freqtext.upper().startswith("G"):                    
                    scale=1e9
                elif freqtext.upper().startswith("T"):
                    scale=1e12
                print "frequency in (%E)Hz\n"%scale
            continue
        data=[float(x) for x in line.split()]
        freq, data = data[0]*scale,data[1:]
        if len(data)!=8:
            print 'wrong line: %s'%line
            continue
        else:
            if AmpArgMode:
                s11,s12,s21,s22 = pol(data[0],data[1]),pol(data[2],data[3]),pol(data[4],data[5]),pol(data[6],data[7])
            else:
                s11,s12,s21,s22 = complex(data[0],data[1]),complex(data[2],data[3]),complex(data[4],data[5]),complex(data[6],data[7])
        if bTwoPorts:           
            a,b,c,d = abcd(s11,s12,s21,s22,z0)
            x1,x2,jwm = rlm(a,b,c,d,z0)
            mat.append((freq,x1,x2,jwm))
        else:
            x1,x2 = z0*(1+s11)/(1-s11), z0*(1+s22)/(1-s22)
            mat.append((freq,x1,x2))
    return array(mat),name

def out2file(mat,filename):
    f = open(filename, 'wb')
    for i in range(len(mat)):
        for j in range(len(mat[0])):
            f.write("%g "%mat[i][j])
        f.write("\r\n")
    f.close()
    

def drawRFData(name, bTwoPorts =True,cord=AutoCord):
    mat,name = read_s2p(name, bTwoPorts =bTwoPorts)
    c = XYChart(830,500)
    c.setPlotArea(50, 35, 700, 420).setGridColor(0xc0c0c0, 0xc0c0c0)
    c.addTitle("Resistance versus Frequency", "timesbi.ttf", 18)
    c.yAxis().setWidth(3)
    c.xAxis().setWidth(3)
    c.xAxis().setTitle("Frequency(Hz)", "arialbi.ttf", 12)
    c.xAxis().setLogScale3()
    c.xAxis().setLabelFormat("{value|E3}")
##    c.yAxis().setLabelFormat("{value|E3}")
##    c.yAxis().setLogScale3()
    colors=[0xff8080,0x8080ff]
    labels=['R1','R2']
    c.yAxis().setTitle("R1")
    c.yAxis().setColors(0xff8080,0xff8080)
    c.yAxis2().setColors(0x8080ff, 0x8080ff)
    c.yAxis2().setTitle("R2")
    c.yAxis2().setWidth(3)
    for r in range(2):
        # print r1 vs freq and r2 vs freq
        layer = c.addLineLayer2()
        freq = mat[:,0]
        layer.setXData(freq.real.tolist())
        temp = layer.addDataSet(mat[:,r+1].real.tolist(),colors[r],labels[r])
        temp.setDataSymbol(SquareSymbol)
        if r>0:
            temp.setUseYAxis2()

    c.addLegend(50, 28, 0, "arialbd.ttf", 10).setBackground(Transparent)
    c.makeChart("%s_RvsF.png"%name)
    mat = array(mat[:,:3].real)
    out2file(mat, "%s_RvsF.txt"%name)
    return mat

def drawLFData(name, bTwoPorts =True,cord=AutoCord):
    mat,name = read_s2p(name,bTwoPorts =bTwoPorts)
    c = XYChart(830,500)
    c.setPlotArea(50, 35, 700, 420).setGridColor(0xc0c0c0, 0xc0c0c0)
    c.addTitle("Inductance versus Frequency", "timesbi.ttf", 18)
    c.yAxis().setWidth(3)
    c.xAxis().setWidth(3)
    c.xAxis().setTitle("Frequency(Hz)", "arialbi.ttf", 12)
    c.xAxis().setLogScale3()
    c.xAxis().setLabelFormat("{value|E3}")
    c.yAxis().setLabelFormat("{value|E2}")
    c.yAxis().setTitle("L1")
    c.yAxis().setColors(0xff8080,0xff8080)
    c.yAxis2().setColors(0x8080ff, 0x8080ff)
    c.yAxis2().setTitle("L2")    
    c.yAxis2().setWidth(3)
    c.yAxis2().setLabelFormat("{value|E2}")
##    c.yAxis().setLogScale3()
    freq = mat[:,0]
    ret = transpose(vstack((freq, mat[:,1].imag/2/pi/freq, mat[:,2].imag/2/pi/freq)))
    colors=[0xff8080,0x8080ff]
    labels=['L1','L2']
    for r in range(2):
        # print l1 vs freq and l2 vs freq
        layer = c.addLineLayer2()
        layer.setXData(freq.real.tolist())
        temp = layer.addDataSet(ret[:,r+1].real.tolist(),colors[r],labels[r])
        temp.setDataSymbol(SquareSymbol)
        if r>0:
            temp.setUseYAxis2()        
    c.addLegend(50, 28, 0, "arialbd.ttf", 10).setBackground(Transparent)
    c.makeChart("%s_LvsF.png"%name)
    out2file(ret, "%s_LvsF.txt"%name)
    return ret

def drawQFData(name, bTwoPorts =True,cord=AutoCord):
    mat,name = read_s2p(name,bTwoPorts =bTwoPorts)
    c = XYChart(830,500)
    c.setPlotArea(50, 35, 700, 420).setGridColor(0xc0c0c0, 0xc0c0c0)
    c.addTitle("Quality Factor(Q) versus Frequency", "timesbi.ttf", 18)
    c.yAxis().setWidth(3)
    c.xAxis().setWidth(3)
    c.xAxis().setTitle("Frequency(Hz)", "arialbi.ttf", 12)
    c.xAxis().setTickDensity(10)
    c.xAxis().setLogScale3()
    c.xAxis().setLabelFormat("{value|E3}")
    c.yAxis().setTitle("Q1")
    c.yAxis().setColors(0xff8080,0xff8080)
    c.yAxis2().setColors(0x8080ff, 0x8080ff)
    c.yAxis2().setTitle("Q2")    
    c.yAxis2().setWidth(3)
    c.yAxis().setLinearScale(-10,100)
    c.yAxis2().setLinearScale(-10,100)
##    c.yAxis().setLogScale3()
    freq = mat[:,0]
    ret = transpose(vstack((freq, mat[:,1].imag/mat[:,1].real, mat[:,2].imag/mat[:,2].real)))
    colors=[0xff8080,0x8080ff]
    labels=['Q1','Q2']
    for r in range(2):
        # print l1 vs freq and l2 vs freq
        layer = c.addLineLayer2()
        layer.setXData(freq.real.tolist())
        temp = layer.addDataSet(ret[:,r+1].real.tolist(),colors[r],labels[r])
        temp.setDataSymbol(SquareSymbol)
        if r>0:
            temp.setUseYAxis2()        
    c.addLegend(50, 28, 0, "arialbd.ttf", 10).setBackground(Transparent)
    c.makeChart("%s_QvsF.png"%name)
    out2file(ret, "%s_QvsF.txt"%name)
    return ret


def drawKFData(name,cord=AutoCord):
    mat,name = read_s2p(name)
    c = XYChart(830,500)
    c.setPlotArea(50, 35, 700, 420).setGridColor(0xc0c0c0, 0xc0c0c0)
    c.addTitle("Coupling Coefficient(k) versus Frequency", "timesbi.ttf", 18)
    c.yAxis().setWidth(3)
    c.xAxis().setWidth(3)
    c.xAxis().setTitle("Frequency(Hz)", "arialbi.ttf", 12)
    c.xAxis().setLabelFormat("{value|E3}")    

    if None == cord.yaxis.min:
        cord.setYmin(0)
    if None == cord.yaxis.max:
        cord.setYmax(1)
    cord.setXMode(True)
    setAxis(c,cord)

#    c.yAxis().setLinearScale3()
    freq = mat[:,0]
    ret = transpose(vstack((freq, abs(mat[:,3].imag/(mat[:,1].imag*mat[:,2].imag)**0.5))))
#    ret = transpose(vstack((freq, (mat[:,3].imag)/2/pi/freq)))
    layer = c.addLineLayer2()
    layer.setXData(freq.real.tolist())
    layer.addDataSet(ret[:,1].real.tolist(),0x80ff80,'K').setDataSymbol(SquareSymbol)
    c.addLegend(50, 28, 0, "arialbd.ttf", 10).setBackground(Transparent)
    c.makeChart("%s_KvsF.png"%name)
    out2file(ret, "%s_KvsF.txt"%name)
    return ret

def drawData(xseq,yseq,xopt="",yopt=""):
    c = XYChart(800,500)
    c.setPlotArea(70, 35, 700, 420).setGridColor(0xc0c0c0, 0xc0c0c0)
    c.addTitle("X versus Y", "timesbi.ttf", 18)
    c.yAxis().setWidth(3)
    c.xAxis().setWidth(3)
    c.xAxis().setTitle("Frequency(Hz)", "arialbi.ttf", 12)

    c.xAxis().setLabelFormat("{value|E3}")    
    c.yAxis().setLabelFormat("{value|E3}")
    if xopt.find("l")!=-1:
        c.xAxis().setLogScale3()
    else:
        c.xAxis().setLinearScale3()
    if yopt.find("l")!=-1:
        c.yAxis().setLogScale3()
    else:
        c.yAxis().setLinearScale3()
    num = len(xseq)
    if num>800:
        freq = array([xseq[int(i*num/800.0)] for i in range(800)])
        ret = array([yseq[int(i*num/800.0)] for i in range(800)])
    else:
        freq = xseq
        ret = yseq

    layer = c.addLineLayer2()
    layer.setXData(freq.real.tolist())
    layer.addDataSet(ret.real.tolist(),0x80ff80,'Y').setDataSymbol(SquareSymbol)
    c.addLegend(50, 28, 0, "arialbd.ttf", 10).setBackground(Transparent)
##    c.makeChart("temp.png")
    return c
    


def drawFile(name, bTwoPorts = False, cord = CurCord):
    if bTwoPorts:
        print 'Generating R,L,K,Q'
        return drawRFData(name,bTwoPorts,cord),drawLFData(name,bTwoPorts,cord),drawKFData(name,cord),drawQFData(name,bTwoPorts,cord)
    else:
        print 'Generating R,L,Q for two individule ports separately'
        return drawRFData(name,bTwoPorts,cord),drawLFData(name,bTwoPorts,cord),drawQFData(name,bTwoPorts,cord)
    
